Amazon Bedrock Knowledge Bases とPineconeで低コストなRAGを構築する

酒井貴央

2024.11.23

先日、Amazon Bedrock Knowledge Basesを使用したRAG（Retrieval-Augmented Generation）の検証を行っていたのですが、予想以上の請求額になっていました。
特にベクトルストアとして使用しているAmazon OpenSearch Serverlessが料金の大半を占めていたため、Freeプランのある Pineconeを使用してRAGを構築してみることにしました。
また、今後の利用頻度が高いことも予想されるため、再利用可能なCloudFormationでの実装を行います。

構成

今回作成するリソースは以下のとおりです。

kb-with-pinecone-20241123

構築のソースコードは下記Githubにアップロードしているので、必要に応じてご参照ください。

ここからはソースコードを一部抜粋しながら解説します。

IAM

template.yaml

  KnowledgeBaseRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub ${AWS::StackName}-role
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: bedrock.amazonaws.com
            Action: sts:AssumeRole
            Condition:
              StringEquals:
                aws:SourceAccount: !Ref AWS::AccountId
              ArnLike:
                AWS:SourceArn: !Sub arn:aws:bedrock:${AWS::Region}:${AWS::AccountId}:knowledge-base/*
      ManagedPolicyArns:
        - !Ref ModelAccessPolicy
        - !Ref RagDataSourceAccessPolicy
        - !Ref SecretsManagerAccessPolicy

  ModelAccessPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Action:
              - bedrock:ListFoundationModels
              - bedrock:ListCustomModels
            Resource: '*'
          - Effect: Allow
            Action:
              - bedrock:InvokeModel
            Resource:
              - !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/${EmbeddingModelId}

  RagDataSourceAccessPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Action:
              - s3:GetObject
              - s3:ListBucket
            Resource:
              - !Sub arn:aws:s3:::${DataSourceBucket}
              - !Sub arn:aws:s3:::${DataSourceBucket}/*
            Condition:
              StringEquals:
                aws:PrincipalAccount: !Ref AWS::AccountId

  SecretsManagerAccessPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Action:
              - secretsmanager:GetSecretValue
            Resource:
              - !Ref PineconeSecretArn

ナレッジベースで使用するIAMロール・ポリシーを作成します。
以下を参考にポリシーは定義しています。

S3

template.yaml

  DataSourceBucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Sub ${AWS::StackName}-rag-datasource-${AWS::AccountId}-${AWS::Region}

データソースとして使用するS3バケットを作成します。

KnowledgeBase

template.yaml

  KnowledgeBase:
    Type: AWS::Bedrock::KnowledgeBase
    Properties:
      KnowledgeBaseConfiguration:
        Type: VECTOR
        VectorKnowledgeBaseConfiguration:
          EmbeddingModelArn: !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/${EmbeddingModelId}
      Name: !Sub ${AWS::StackName}-knowledgebase
      RoleArn: !GetAtt KnowledgeBaseRole.Arn
      StorageConfiguration:
        Type: PINECONE
        PineconeConfiguration:
          ConnectionString: !Ref PineconeEndpoint
          CredentialsSecretArn: !Ref PineconeSecretArn
          FieldMapping:
            MetadataField: metadata
            TextField: text

  S3DataSource:
    Type: AWS::Bedrock::DataSource
    Properties:
      DataSourceConfiguration:
        Type: S3
        S3Configuration:
          BucketArn: !GetAtt DataSourceBucket.Arn
      Name: !Sub ${AWS::StackName}-data-source
      KnowledgeBaseId: !Ref KnowledgeBase

ナレッジベースを作成し、作成したS3バケットをデータソースとして紐付けます。

Pinecone Index

createIndex.js

// Import the Pinecone library
import { Pinecone } from '@pinecone-database/pinecone';

// Initialize a Pinecone client with your API key
const pc = new Pinecone({ apiKey: 'YOUR_API_KEY' });

// Create a serverless index
const indexName = "example-index"

await pc.createIndex({
  name: indexName,
  dimension: 1024,
  metric: 'cosine',
  spec: { 
    serverless: { 
      cloud: 'aws', 
      region: 'us-east-1' 
    }
  } 
}); 

// Describe the endpoint of the created index
const indexInfo = await pc.describeIndex(indexName);
const endpoint = indexInfo.host;

console.log(`PineconeEndPoint: https://${endpoint}`);

公式ドキュメントを参考に、SDKでインデックスの作成を行います。
使用する埋め込みモデルによって、dimensionの値を変える必要があるので注意してください。
今回は Titan Text Embeddings v2 を使用するので、1024としています。
CFnテンプレート側でPinecone インデックス管理ページのエンドポイント URLが必要になりますので、コンソールログへのURL出力も行なっています。

構築

ここからは実際に構築を行います。
前提として以下の準備が必要なので、適宜実施してください。

AWS認証情報の設定
- AWS CLIの設定と認証情報ファイルの設定
Bedrock モデルの有効化
- Amazon Bedrock 基盤モデルへのアクセスの追加または削除
Pineconeのアカウント及びAPIキー作成
- APIキーを取得する
createIndex.js
- YOUR_API_KEY: PineconeのAPIキーを設定してください。
- example-index: Pineconeのインデックス名を設定してください。

その他にJSの実行環境も必要になりますが、ご自身の状況に合わせて対応を行なってください。

GitHubからリポジトリをクローンし、ディレクトリを移動します。

$ git clone https://github.com/sakai-classmethod/knowledgebase-with-pinecone.git
$ cd knowledgebase-with-pinecone

必要なパッケージをインストールします。(npm等でも可)

$ bun install

Pineconeのアカウントを作成した後に、createIndex.js内のYOUR_API_KEYとexample-indexを実際に使用する値に書き換えてスクリプトを実行します。

$ bun createIndex.js
PineconeEndPoint: https://xxxxxx.pinecone.io

PineconeのAPIキーを格納するためのSecrets Managerを作成します。
今回はオレゴンリージョンを指定します。

$ aws secretsmanager create-secret --region us-west-2 --name PineconeSecret --secret-string "{\"apiKey\":\"YOUR_API_KEY\"}"
{
    "ARN": "arn:aws:secretsmanager:us-west-2:xxx:secret:PineconeSecret-xxx",
    "Name": "PineconeSecret",
    "VersionId": "xxx-xxx-xxx-xxx-xxx"
}

最後にCloudFormationをデプロイします。（コンソールやAWS CLIからのデプロイでも可）
同じくオレゴンリージョンを指定します。

$ rain deploy -r us-west-2 template.yaml knowledgebase-stack
Enter a value for parameter 'EmbeddingModelId' (default value: amazon.titan-embed-text-v2:0): <変更がなければそのままEnter>
Enter a value for parameter 'PineconeEndPoint': <Pineconeインデックス管理ページのエンドポイントURL>
Enter a value for parameter 'PineconeSecretArn': <PineconeのAPIキーを格納するためのSecrets ManagerのArn>
CloudFormation will make the following changes:
Stack knowledgebase-stack:
  + AWS::S3::Bucket DataSourceBucket
  + AWS::IAM::Role KnowledgeBaseRole
  + AWS::Bedrock::KnowledgeBase KnowledgeBase
  + AWS::IAM::ManagedPolicy ModelAccessPolicy
  + AWS::IAM::ManagedPolicy RagDataSourceAccessPolicy
  + AWS::Bedrock::DataSource S3DataSource
  + AWS::IAM::ManagedPolicy SecretsManagerAccessPolicy
Do you wish to continue? (Y/n) 
Deploying template 'template.yaml' as stack 'knowledgebase-stack' in us-west-2.
Stack knowledgebase-stack: CREATE_COMPLETE
  Outputs:
    KnowledgeBaseId: xxxxxxxxxx

動作確認

作成されたS3バケットに任意のファイルをアップロードして、データソースを同期します。

upload-data-20241123

sync-data-source-20241123

テキストモデルを選択してテストしてみると問題なく動作していることが確認できました。

q1-kb-20241123

データソースに追加されていないことについても確認してみましたが、問題なさそうです。

q2-kb-20241123

最後にPineconeコンソールからインデックスのレコードを確認してみます。

pinecone-index-20241123

同期したデータがベクトルデータとして登録されていることが、確認できました。

まとめ

今回はAWS利用費の高騰をきっかけに、ベクトルデータベースのPineconeを使用したRAGを構築してみました。
Amazon OpenSearch Serverlessでしか使えない機能はあるものの、ちょっとした検証やコストを抑えたい場面ではPineconeの使い勝手がいいことがわかりました。
AWS Marketplaceからの利用も可能なので、請求を一元化したい場合は以下ドキュメントを参考に使ってみてください。
どなたかの参考になれば幸いです。